from kuai.sim import *
from kuai.mol.io import read_molfile
from kuai.sim.dff2kuai import read_ppffile
from kuai.unit import default_units
from numpy import zeros, double, sum, random
from scipy import optimize 
from time import clock

class EnergyFunction:
    def __init__(self, model, index):
        self.model = model
        self.index = index
        self.job = SimulationJob()
        
    def __call__(self, x):
        self.model.coords = x
        funcs = make_efunc(self.model, self.index)
        e = self.job.getE(self.model, funcs)
        result = sum(e)
        return result

        
class GradientFunction:
    def __init__(self, model, index):
        self.model = model
        self.index = index
        self.job = SimulationJob()
        
    def __call__(self, x):
        self.model.coords = x
        funcs = make_efunc(self.model, self.index)
        _, g = self.job.getEG(model, funcs)
        return g

if __name__ == '__main__':
    import sys
    mol = read_molfile(sys.argv[1])
    index, parameters = read_ppffile(sys.argv[2])
    model = setup_model(mol, index, parameters)
    natoms = len(mol.atoms)

    efunc = EnergyFunction(model, index)
    gfunc = GradientFunction(model, index)

    random.seed(3)
    startX =random.rand(natoms*3) * 10
    result = optimize.fmin_bfgs(efunc, startX, gfunc, full_output=True)
    
    fe = default_units.format(1, 'kcal/mol')[0] # Factor of Energy Unit
    print result[1] * fe
    
    with open(sys.argv[3], "wt") as file:
        file.write("%d\n Optimized Molecule\n" % natoms)
        for i in range(natoms):
            file.write("%s %10.4f %10.4f %10.4f\n" % (mol.atoms[i].symbol, result[0][i*3], result[0][i*3+1], result[0][i*3+2]))
